package org.zjvis.datascience.service.csv.postgresql;

import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.NoArgsConstructor;
import lombok.Singular;
import org.postgresql.copy.CopyManager;
import org.postgresql.core.BaseConnection;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.List;
import java.util.stream.Collectors;

@Builder
@AllArgsConstructor
@NoArgsConstructor
public class PostgresqlCsvTransformer {
    private String host;
    private int port;
    private String username;
    private String password;
    private File csvFile;
    private String databaseName;
    // TODO use interface
    private PostgresqlTableCreator tableCreator;
    @Builder.Default
    private boolean dropTableIfExists = true;
    @Builder.Default
    private boolean createTableIfNotExists = true;
    @Builder.Default
    private String driverName = "org.postgresql.Driver";
    @Singular("columnName")
    private List<String> columnNameList;

    public long csvToTable() throws SQLException, IOException, ClassNotFoundException {
        Class.forName(driverName);
        try (Connection connection = DriverManager.getConnection(connectionUrl(), username, password)) {
            if (dropTableIfExists) {
                try (PreparedStatement statement
                             = connection.prepareStatement(tableCreator.dropTableIfExistsSql())) {
                    statement.execute();
                }
            }
            if (createTableIfNotExists) {
                try (PreparedStatement statement
                             = connection.prepareStatement(tableCreator.createTableIfNotExistsSql())) {
                    statement.execute();
                }
            }
            String columnDefinition = columnNameList.isEmpty()
                    ? ""
                    : String.format("(%s)", columnNameList.stream().map(columnName -> String.format("\"%s\"", columnName)).collect(Collectors.joining(", ")));
            return new CopyManager((BaseConnection) connection)
                    .copyIn(
                            String.format("COPY %s%s FROM STDIN (FORMAT csv, HEADER)",
                                    tableCreator.getTableName(), columnDefinition),
                            new BufferedReader(new FileReader(csvFile))
                    );
        }
    }

    private String connectionUrl() {
        return String.format("jdbc:postgresql://%s:%s/%s", host, port, databaseName);
    }
}
